package pl.mmadej.ds4java.ds.tree.binary.avl;


import pl.mmadej.ds4java.ds.hash.ArrayHashMap;
import pl.mmadej.ds4java.ds.hash.HashMap;
import pl.mmadej.ds4java.ds.tree.binary.BinaryTreeNode;
import pl.mmadej.ds4java.ds.tree.binary.bst.BinarySearchTree;


public class AVLTree<T extends Comparable<T>> extends BinarySearchTree<T> {

    //augmentation - adding height to node
    private HashMap<BinaryTreeNode<T>, Integer> nodeHeight = new ArrayHashMap<BinaryTreeNode<T>, Integer>(100, 0.75);

    @Override
    protected void afterNodeAdded(BinaryTreeNode<T> newNode) {
        super.afterNodeAdded(newNode);
        updateNodeHeightWithChildrenMaxHeightPlusOne(newNode);
        balanceTreeStartingFrom(newNode.getParent());
    }

    @Override
    protected void afterNodeDeleted(BinaryTreeNode<T> lowestAffectedParent) {
        super.afterNodeDeleted(lowestAffectedParent);
        balanceTreeStartingFrom(lowestAffectedParent);
    }

    @Override
    protected void afterNodeRotatedLeft(BinaryTreeNode<T> nodeToGoDownOnLeftSide, BinaryTreeNode<T> rightChildOfNode) {
        super.afterNodeRotatedLeft(nodeToGoDownOnLeftSide, rightChildOfNode);
        updateNodeHeightWithChildrenMaxHeightPlusOne(nodeToGoDownOnLeftSide);
        updateNodeHeightWithChildrenMaxHeightPlusOne(rightChildOfNode);
    }

    @Override
    protected void afterNodeRotatedRight(BinaryTreeNode<T> nodeToGoDownOnRightSide, BinaryTreeNode<T> leftChildOfNode) {
        super.afterNodeRotatedRight(nodeToGoDownOnRightSide, leftChildOfNode);
        updateNodeHeightWithChildrenMaxHeightPlusOne(nodeToGoDownOnRightSide);
        updateNodeHeightWithChildrenMaxHeightPlusOne(leftChildOfNode);
    }

    private void balanceTreeStartingFrom(BinaryTreeNode<T> parent) {
        while (parent != null) {
            updateNodeHeightWithChildrenMaxHeightPlusOne(parent);
            if (rotationRequired(parent)) {
                if (parentShouldSingleRotateWithLeftChild(parent)) {
                    rotateRight(parent);
                } else if (parentShouldSingleRotateWithRightChild(parent)) {
                    rotateLeft(parent);
                } else if (parentShouldDoubleRotateWithLeftChildAssumingThatParentDidNotSingleRotateWithLeftChild(parent)) {
                    // double rotate with left child
                    rotateLeft(parent.getLeftChild());
                    rotateRight(parent);
                } else {
                    // double rotate with right child
                    rotateRight(parent.getRightChild());
                    rotateLeft(parent);
                }
            }
            parent = parent.getParent();
        }
    }

    private void updateNodeHeightWithChildrenMaxHeightPlusOne(BinaryTreeNode<T> newNode) {
        updateNodeHeight(newNode, maxHeightFromChildren(newNode) + 1);
    }

    private void updateNodeHeight(BinaryTreeNode<T> node, int height) {
        nodeHeight.insert(node, height);
    }

    private Integer maxHeightFromChildren(BinaryTreeNode<T> parent) {
        return rightChildHeight(parent) > leftChildHeight(parent) ?
                rightChildHeight(parent) : leftChildHeight(parent);
    }

    private boolean parentShouldSingleRotateWithLeftChild(BinaryTreeNode<T> parent) {
        return leftChildHeight(parent) == rightChildHeight(parent) + 2
                && leftChildHeight(parent.getLeftChild()) == rightChildHeight(parent) + 1;
    }

    private boolean parentShouldSingleRotateWithRightChild(BinaryTreeNode<T> parent) {
        return rightChildHeight(parent) == leftChildHeight(parent) + 2
                && rightChildHeight(parent.getRightChild()) == leftChildHeight(parent) + 1;
    }

    private boolean parentShouldDoubleRotateWithLeftChildAssumingThatParentDidNotSingleRotateWithLeftChild(BinaryTreeNode<T> parent) {
        return leftChildHeight(parent) == rightChildHeight(parent) + 2 ;
            //    && rightChildHeight(parent.getLeftChild()) == rightChildHeight(parent) + 1;

    }

    private boolean rotationRequired(BinaryTreeNode<T> parent) {
        return Math.abs(rightChildHeight(parent) - leftChildHeight(parent)) > 1;
    }

    private int rightChildHeight(BinaryTreeNode<T> parent) {
        return parent.getRightChild() != null ? nodeHeight.search(parent.getRightChild()) : 0;
    }

    private int leftChildHeight(BinaryTreeNode<T> parent) {
        return parent.getLeftChild() != null ? nodeHeight.search(parent.getLeftChild()) : 0;
    }
}
